/**
 * Copyright (c) Huawei Technologies Co., Ltd. 2023. All rights reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 * http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#ifndef AIR_CXX_RUNTIME_V2_OP_IMPL_DYNAMIC_UNSORTED_SEGMENT_SUM_H_
#define AIR_CXX_RUNTIME_V2_OP_IMPL_DYNAMIC_UNSORTED_SEGMENT_SUM_H_

#include <cstdint>
#include <vector>
#include "op_tiling.h"

namespace TbeReduce {
    const std::string UNSORTED_SEGMENT_SUM_OP_TYPE = "UnsortedSegmentSum";
    const int64_t CONST_IDX = 2;
    const int32_t BYTE_BLOCK = 32;
    const int32_t MIN_ELE_SIZE_USING_ALL_CORE = 1024;
    const int32_t MASK_FP32 = 64;
    const int32_t MASK_INT32 = 64;
    const int32_t MASK_FP16 = 128;
    const int32_t MAX_REPEAT_TIME = 255;
    const int32_t FP32_ELE_NUM_ALIGN_32B = 8;
    const int32_t BYTE_FULL_MASK = 256;
    const int32_t MULTI = 4;
    const int32_t FP16_BLOCK_NUM = 16;
    const int32_t INT32_BLOCK_NUM = 8;
    // dtype
    const std::string DTYPE_FP32 = "float32";
    const std::string DTYPE_FP16 = "float16";
    const std::string DTYPE_INT32 = "int32";
    const std::string DTYPE_INT8 = "int8";
    const std::string DTYPE_UINT8 = "uint8";

    // ub tensor num
    const int32_t UB_TENSOR_NUM_FP32_INPUT_LAST_AXIS_ALIGN = 2;
    const int32_t UB_TENSOR_NUM_FP32_INPUT_LAST_AXIS_ONE = 6;
    const int32_t UB_TENSOR_NUM_FP32_INPUT_ONE_DIM = 3;
    const int32_t UB_TENSOR_NUM_FP32_PAD = 6;
    const int32_t PAD_ONCE = 128;
    const int32_t PAD_LASE_AXIS = 168;
    const int32_t UB_TENSOR_NUM_FP32_INPUT_LAST_AXIS_NOT_ALIGN = 3;
    const int32_t UB_TENSOR_NUM_FP32_INPUT_LAST_AXIS_NOT_ALIGN_HP = 3;

    // fp32 select key
    const int32_t SELECT_KEY_MODE_FP32_INPUT_LAST_AXIS_ALIGN_SMALL_E = 1;
    const int32_t SELECT_KEY_MODE_FP32_INPUT_LAST_AXIS_ONE = 2;
    const int32_t SELECT_KEY_MODE_FP32_INPUT_ONE_DIM = 2;
    const int32_t SELECT_KEY_MODE_FP32_INPUT_LAST_AXIS_NOT_ALIGN_SMALL_E = 4;
    const int32_t SELECT_KEY_MODE_FP32_INPUT_LAST_AXIS_ALIGN_BIG_E = 5;
    const int32_t SELECT_KEY_MODE_FP32_INPUT_LAST_AXIS_NOT_ALIGN_BIG_E = 6;
    const int32_t SELECT_KEY_MODE_FP32_INPUT_LAST_AXIS_ONE_MODIFY = 7;
    const int32_t SELECT_KEY_MODE_FP32_INPUT_LAST_AXIS_ONE_MULTI = 8;
    const int32_t SELECT_KEY_MODE_FP32_INPUT_NUM_SEGMENT_ONE = 17;
    const int32_t SELECT_KEY_MODE_FP32_INPUT_LAST_AXIS_ALIGN_SMALL_E_HP = 18;
    const int32_t SELECT_KEY_MODE_FP32_INPUT_LAST_AXIS_NOT_ALIGN_SMALL_E_HP = 19;
    const int32_t SELECT_KEY_MODE_FP32_INPUT_LAST_AXIS_NOT_ALIGN_SMALL_E_HP_PAD = 20;

    // int32 select key
    const int32_t SELECT_KEY_MODE_NO_ATOMIC_SMALL_E_SMALL_ID = 9;
    const int32_t SELECT_KEY_MODE_NO_ATOMIC_SMALL_E_BIG_ID = 10;
    const int32_t SELECT_KEY_MODE_NO_ATOMIC_BIG_E_SMALL_ID = 11;
    const int32_t SELECT_KEY_MODE_NO_ATOMIC_BIG_E_BIG_ID = 12;
    const int32_t SELECT_KEY_MODE_NO_ATOMIC_SMALL_E_SMALL_ID_SMALLBLOCK = 13;
    const int32_t SELECT_KEY_MODE_NO_ATOMIC_SMALL_E_BIG_ID_SMALLBLOCK = 14;
    const int32_t SELECT_KEY_MODE_NO_ATOMIC_NUM_SEGMENT_ONE = 15;
    const int32_t SELECT_KEY_MODE_NO_ATOMIC_ALL_IN_ALIGN = 16;
    enum EleByte {
        FP16_BYTE = 2, FP32_BYTE = 4, INT32_BYTE = 4, INT64_BYTE = 8, INT8_BYTE = 1, UINT8_BYTE = 1
    };

    static const std::vector <std::string> COMPILE_INFO_KEY = {"core_num", "ub_size", "ub_tensor_num", "impl_mode"};
    const int32_t PARA_CORE_NUM_0 = 0;
    const int32_t PARA_UB_SIZE_1 = 1;
    const int32_t PARA_UB_TENSOR_NUM_2 = 2;
    const int32_t PARA_IMPL_MODE_3 = 3;
    const int32_t UB_TENSOR_NUM_2 = 2;

    const int32_t UB_TENSOR_MAX_X_NUM = 16000;

    struct CompileParas {
        // std::shared_ptr<AutoTilingHandler> outer_compile_info;
        int32_t core_num;
        int32_t ub_size;
        int32_t ub_tensor_num;
        int32_t impl_mode;
        bool is_tik{false};
    };

    struct CommParas {
        int32_t e_num;
        int32_t x_e_num;
        int32_t num_segments;
        int32_t ub_size;
        int32_t impl_mode;
        int32_t ub_tensor_ele_num_pad;
        int32_t e_size_align;
        bool is_pad;
    };

    struct TilingParamsFp32 {
        // common params
        int32_t select_key_input_scalar;
        int32_t need_core_num_input_scalar;

        // input data params
        // front core
        int32_t input_ele_num_front_core_input_scalar;
        // front part front core
        int32_t input_mov_times_gm2ub_front_part_front_core_input_scalar;
        int32_t input_front_burst_len_front_part_front_core_input_scalar;
        int32_t input_last_burst_len_front_part_front_core_input_scalar;
        int32_t input_front_ele_num_ub_front_part_front_core_input_scalar;
        int32_t input_last_ele_num_ub_front_part_front_core_input_scalar;
        int32_t input_front_rows_front_part_front_core_input_scalar;
        int32_t input_last_rows_front_part_front_core_input_scalar;
        // last part front core
        int32_t input_mov_times_gm2ub_last_part_front_core_input_scalar;
        int32_t input_front_burst_len_last_part_front_core_input_scalar;
        int32_t input_last_burst_len_last_part_front_core_input_scalar;
        int32_t input_front_ele_num_ub_last_part_front_core_input_scalar;
        int32_t input_last_ele_num_ub_last_part_front_core_input_scalar;
        int32_t input_front_rows_last_part_front_core_input_scalar;
        int32_t input_last_rows_last_part_front_core_input_scalar;
        // last core
        int32_t input_ele_num_last_core_input_scalar;
        // front part last core
        int32_t input_mov_times_gm2ub_front_part_last_core_input_scalar;
        int32_t input_front_burst_len_front_part_last_core_input_scalar;
        int32_t input_last_burst_len_front_part_last_core_input_scalar;
        int32_t input_front_ele_num_ub_front_part_last_core_input_scalar;
        int32_t input_last_ele_num_ub_front_part_last_core_input_scalar;
        int32_t input_front_rows_front_part_last_core_input_scalar;
        int32_t input_last_rows_front_part_last_core_input_scalar;
        // last part last core
        int32_t input_mov_times_gm2ub_last_part_last_core_input_scalar;
        int32_t input_front_burst_len_last_part_last_core_input_scalar;
        int32_t input_last_burst_len_last_part_last_core_input_scalar;
        int32_t input_front_ele_num_ub_last_part_last_core_input_scalar;
        int32_t input_last_ele_num_ub_last_part_last_core_input_scalar;
        int32_t input_front_rows_last_part_last_core_input_scalar;
        int32_t input_last_rows_last_part_last_core_input_scalar;

        // e num params
        int32_t e_num_input_scalar;
        int32_t e_mov_times_gm2ub_input_scalar;
        int32_t e_ub2gm_front_burst_len_input_scalar;
        int32_t e_num_front_part_input_scalar;
        int32_t e_ub2gm_last_burst_len_input_scalar;
        int32_t e_gm2ub_last_burst_len_input_scalar;
        int32_t e_num_last_part_input_scalar;

        // ids params
        int32_t ids_size_input_scalar;
        int32_t ids_ele_num_front_core_input_scalar;
        int32_t ids_mov_times_gm2ub_front_core_input_scalar;
        int32_t ids_front_burst_len_front_core_input_scalar;
        int32_t ids_last_burst_len_front_core_input_scalar;
        int32_t ids_ele_num_ub_front_part_front_core_input_scalar;
        int32_t ids_ele_num_ub_last_part_front_core_input_scalar;
        int32_t ids_ele_num_last_core_input_scalar;
        int32_t ids_mov_times_gm2ub_last_core_input_scalar;
        int32_t ids_front_burst_len_last_core_input_scalar;
        int32_t ids_last_burst_len_last_core_input_scalar;
        int32_t ids_ele_num_ub_front_part_last_core_input_scalar;
        int32_t ids_ele_num_ub_last_part_last_core_input_scalar;

        // output init params
        int32_t output_ub_init_last_repeat_time_front_part_front_core_input_scalar;
        int32_t output_ub_init_times_front_part_front_core_input_scalar;
        int32_t output_ub_init_last_repeat_time_last_part_front_core_input_scalar;
        int32_t output_ub_init_times_last_part_front_core_input_scalar;
        int32_t output_ub_init_last_repeat_time_front_part_last_core_input_scalar;
        int32_t output_ub_init_times_front_part_last_core_input_scalar;
        int32_t output_ub_init_last_repeat_time_last_part_last_core_input_scalar;
        int32_t output_ub_init_times_last_part_last_core_input_scalar;
        int32_t input_last_axis_align_front_part_ele_num_input_scalar;
        int32_t input_last_axis_align_floor_ele_num_input_scalar;
        int32_t last_part_vadd_mask_input_scalar;
        int32_t output_ub_init_last_row_last_repeat_time_front_part_front_core_input_scalar;
        int32_t output_ub_init_last_row_times_front_part_front_core_input_scalar;
        int32_t output_ub_init_last_row_last_repeat_time_last_part_front_core_input_scalar;
        int32_t output_ub_init_last_row_times_last_part_front_core_input_scalar;
        int32_t output_ub_init_last_row_last_repeat_time_front_part_last_core_input_scalar;
        int32_t output_ub_init_last_row_times_front_part_last_core_input_scalar;
        int32_t output_ub_init_last_row_last_repeat_time_last_part_last_core_input_scalar;
        int32_t output_ub_init_last_row_times_last_part_last_core_input_scalar;
        int32_t num_segments;
        int32_t tiling_core_num;
        int32_t repeat_front_front_part_front_core;
        int32_t col_sub_block_front_front_part_front_core;
        int32_t repeat_last_front_part_front_core;
        int32_t col_sub_block_last_front_part_front_core;
        int32_t repeat_front_last_part_front_core;
        int32_t col_sub_block_front_last_part_front_core;
        int32_t repeat_last_last_part_front_core;
        int32_t col_sub_block_last_last_part_front_core;
        int32_t repeat_front_front_part_last_core;
        int32_t col_sub_block_front_front_part_last_core;
        int32_t repeat_last_front_part_last_core;
        int32_t col_sub_block_last_front_part_last_core;
        int32_t repeat_front_last_part_last_core;
        int32_t col_sub_block_front_last_part_last_core;
        int32_t repeat_last_last_part_last_core;
        int32_t col_sub_block_last_last_part_last_core;
        int32_t e_num_sub;
        int32_t vadd_repeat_255;
        int32_t vadd_repeat_64;
        int32_t vadd_repeat_last;
        int32_t move_pad;
        int32_t max_cache_n_num;
        int32_t repeat_remove_pad;
        int32_t col_block_remove_pad;
        int32_t cache_num_block;
    };

    struct TilingParamsInt32 {
        // common params
        int32_t select_key_input_scalar;
        int32_t need_core_num_input_scalar;
        int32_t num_segments_front_core_input_scalar;
        int32_t num_segments_last_core_input_scalar;

        // ids params
        int32_t ids_size_input_scalar;
        int32_t ids_mov_times_gm2ub_input_scalar;
        int32_t ids_ele_num_ub_front_part_input_scalar;
        int32_t ids_front_burst_len_input_scalar;
        int32_t ids_ele_num_ub_last_part_input_scalar;
        int32_t ids_last_burst_len_input_scalar;

        // e num params
        int32_t e_num_input_scalar;
        int32_t e_mov_times_gm2ub_input_scalar;
        int32_t e_ub2gm_front_burst_len_input_scalar;
        int32_t e_num_front_part_input_scalar;
        int32_t repeat_time_front_part_input_scalar;
        int32_t e_ub2gm_last_burst_len_input_scalar;
        int32_t e_num_last_part_input_scalar;
        int32_t repeat_time_last_part_input_scalar;
        int32_t align_scalar;
        int32_t align_scalar_lastcore;

        int32_t e_gm2ub_front_burst_len_input_scalar;
        int32_t e_gm2ub_last_burst_len_input_scalar;
        int32_t num_segment_max;
        int32_t num_segment_max_time;
        int32_t num_segment_max_time_lastcore;
        int32_t front_num_segment;
        int32_t front_num_segment_last;
        int32_t front_num_segment_lastcore;
        int32_t front_num_segment_last_lastcore;
        int32_t e_ub2gm_front_burst_len_input_scalar_lastcore;
        int32_t e_ub2gm_last_burst_len_input_scalar_lastcore;
        int32_t repeat_times;
        int32_t repeat_times_last_part;
        int32_t repeat_times_last_part_lastcore;
        int32_t e_mov_times_gm2ub_input_scalar_lastcore;
        int32_t repeat_time_front_part_input_scalar_lastcore;
        int32_t num_segments;

        int32_t tiling_core_num;
    };

    bool SegmentTIKTiling(const std::string &op_type, const TeOpParas &opParas, const CompileParas &op_info,
                          OpRunInfo &run_info);
    int32_t ComputeDivRemainders(const int32_t &num, const int32_t &factor, const int32_t &times);
    void Div(const int32_t &num, const int32_t &factor, int32_t &front_part_num, int32_t &last_part);
    int32_t UssCeil(const int32_t &num, const int32_t &factor);
    int32_t UssCeilDiv(const int32_t &num, const int32_t &factor);
    int32_t UssCeilDivNoAtomic(const int32_t &num, const int32_t &factor, const int32_t &e_size,
        const int32_t &output_ub_ele_num_one_row);
    bool GetTilingMode(const CommParas &commParas, const int32_t &x_ub_size, int32_t &selectKey);
    bool GetTilingModeNoAtomic(
        std::vector <int64_t> &input_shape, const int32_t &e_size, const int32_t &ids_size,
        const std::string &input_dtype, const std::string &ids_dtype, const int32_t &ub_tensor_size,
        const int32_t &ub_tensor_size_input, int32_t &select_key, int32_t &e_once_num, int32_t &id_once_num,
        const int32_t &need_core, const int32_t &output_ub_ele_num_one_row, int32_t &num_segment_max,
        const int32_t &mask, const int32_t &all_size, int32_t &num_segments);
    bool GetEleDtype(const std::string dtype, EleByte &elebyte);
    bool IsUsingAllCore(const int32_t &ids_size, const int32_t &core_num, int32_t &need_core_num, const int32_t &e_size,
        const int32_t &num_segments);
    bool IsUsingAllCoreByNumSegments(const int32_t &num_segments, const int32_t &core_num,
        int32_t &need_core_num, const int32_t &e_size, int32_t &output_ub_ele_num_one_row);
    void ComputeUbTensorSizeNoAtomic(
        const int32_t &ub_size, std::vector <int64_t> &input_shape,
        const std::string &input_dtype, const int32_t &e_size, int32_t &ub_tensor_size_id,
        int32_t &ub_tensor_size_input,
        int32_t &ub_tensor_size_output, const int32_t &output_ub_ele_num_one_row,
        const int32_t &need_core_num, const int32_t &mask, int32_t num_segments);
    void NumSegmentOne(
        int32_t &e_mov_times_gm2ub_input_scalar, const int32_t &max_ele_num_one_ub_tensor,
        int32_t &e_num_front_part_input_scalar, int32_t &e_ub2gm_front_burst_len_input_scalar,
        int32_t &repeat_times, int32_t &repeat_time_front_part_input_scalar, int32_t &e_num_last_part_input_scalar,
        int32_t &e_ub2gm_last_burst_len_input_scalar, int32_t &repeat_times_last_part, int32_t &mask,
        const EleByte &ele_byte, int32_t &num_segments_front_core_input_scalar,
        int32_t &repeat_time_last_part_input_scalar);
    void ComputeUbTensorSize(CommParas &commParas, int32_t &x_ub_size, int32_t &ids_ub_size,
        int32_t &x_ub_size_pad);
    void ComputeEleNumOneCore(const int32_t &min_ele_num, const int32_t &ids_num, const int32_t &core_num,
        const int32_t &e_size, int32_t &ids_ele_num_front_core, int32_t &ids_ele_num_last_core,
        int32_t &input_ele_num_front_core, int32_t &input_ele_num_last_core,
        const int32_t &num_segments);
    void ComputeInputParamsMovGm2ub(const int32_t &ub_tensor_size, const EleByte &input_ele_byte,
        const int32_t &input_ele_num, const int32_t &e_size,
        const int32_t &e_mov_times_gm2ub,
        const int32_t &ids_ele_num_ub, int32_t &input_mov_times_gm2ub,
        int32_t &input_front_burst_len, int32_t &input_last_burst_len,
        int32_t &input_ele_num_ub_front_part, int32_t &input_ele_num_ub_last_part,
        int32_t &input_rows_front_part, int32_t &input_rows_last_part);
    void ComputeIdsParamsMovGm2ub(const int32_t &ids_ele_num_one_core, const int32_t &ub_tensor_size,
        const EleByte &ids_ele_byte, int32_t &ids_mov_times_gm2ub,
        int32_t &ids_front_burst_len,
        int32_t &ids_last_burst_len, int32_t &ids_ele_num_ub_front_part,
        int32_t &ids_ele_num_ub_last_part);
    void ComputeIdsParamsMovGm2ubNoAtomic(const int32_t &ids_ele_num_one_core, const int32_t &id_once_num,
        const EleByte &ids_ele_byte, int32_t &ids_mov_times_gm2ub,
        int32_t &ids_front_burst_len, int32_t &ids_last_burst_len,
        int32_t &ids_ele_num_ub_front_part, int32_t &ids_ele_num_ub_last_part);
    void ComputeInitOutputUbParams(const int32_t &ids_ele_num, const int32_t &output_ub_ele_num_one_row,
        int32_t &output_ub_init_last_repeat_time, int32_t &output_ub_init_times,
        int32_t &repeat_pad, int32_t &sub_block, const CommParas &comm_paras);
    void ComputeNumSegmentsParams(
        const int32_t &need_core_num, const int32_t &num_segmens, int32_t &num_segmens_front_core,
        int32_t &num_segmens_last_core, const int32_t &e_size, int32_t &output_ub_ele_num_one_row);
    void ComputeENumParams(
        const std::string &input_dytpe, const int32_t &e_num, const EleByte &ele_byte,
        const int32_t &e_once_num, int32_t &e_mov_times_gm2ub_input_scalar,
        int32_t &e_ub2gm_front_burst_len_input_scalar,
        int32_t &e_num_front_part_input_scalar, int32_t &repeat_time_front_part_input_scalar,
        int32_t &e_ub2gm_last_burst_len_input_scalar, int32_t &e_num_last_part_input_scalar,
        int32_t &repeat_time_last_part_input_scalar, int32_t &align_scalar, int32_t &align_scalar_lastcore,
        int32_t &e_gm2ub_front_burst_len_input_scalar, int32_t &e_gm2ub_last_burst_len_input_scalar,
        int32_t &num_segments_front_core_input_scalar, int32_t &num_segments_last_core_input_scalar,
        const int32_t &need_core, int32_t &num_segment_max, int32_t &num_segment_max_time,
        int32_t &num_segment_max_time_lastcore, int32_t &front_num_segment, int32_t &front_num_segment_last,
        int32_t &front_num_segment_lastcore, int32_t &front_num_segment_last_lastcore,
        int32_t &e_ub2gm_front_burst_len_input_scalar_lastcore,
        int32_t &e_ub2gm_last_burst_len_input_scalar_lastcore,
        const int32_t &all_size, const int32_t &num_segments, int32_t &repeat_times,
        int32_t &repeat_times_last_part,
        int32_t &repeat_times_last_part_lastcore, int32_t &e_mov_times_gm2ub_input_scalar_lastcore,
        int32_t &repeat_time_front_part_input_scalar_lastcore);
    void InitTilingParams(TilingParamsFp32 &params, int32_t num_segments);
    void InitTilingParams(TilingParamsInt32 &params, int32_t num_segments);
    void WriteTilingParams(const std::string &opType, const TilingParamsFp32 &params, OpRunInfo &run_info);
    void WriteTilingParams(const std::string &opType, const TilingParamsInt32 &params, OpRunInfo &run_info);
    void PrintTilingParams(const std::string &op_type, const TilingParamsFp32 &params);
    void PrintTilingParams(const std::string &op_type, const TilingParamsInt32 &params);

}  // namespace optiling
#endif  // AIR_CXX_RUNTIME_V2_OP_IMPL_DYNAMIC_ATOMIC_ADDR_CLEAN_IMPL_H_
